package com.nacol.demo.gateway.loadbalance;

import cn.hutool.core.util.StrUtil;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;
import java.util.stream.Collectors;

/**
 * Description for this class.
 *
 * @Author : kimmking(kimmking@apache.org)
 * @create 2023/12/13 02:19
 */
public class NearlyLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private final String serviceId;
    private final DiscoveryClient discoveryClient;
    private final Random random = new Random();

    private final String dc = System.getProperty("dc", "A");

    public NearlyLoadBalancer(String serviceId, LoadBalancerClientFactory loadBalancerClientFactory) {
        this.serviceId = serviceId;
        this.discoveryClient = loadBalancerClientFactory.getLazyProvider(serviceId, DiscoveryClient.class).getIfAvailable();
    }

    public Mono<Response<ServiceInstance>> choose(Request request) {
        List<ServiceInstance> instances = discoveryClient.getInstances(serviceId);
        if (instances.isEmpty()) {
            return Mono.empty();
        }

        List<ServiceInstance> nearlyInstances = instances.stream()
                .filter(instance -> StrUtil.equals(instance.getMetadata().get("dc"), dc))
                .collect(Collectors.toList());

        if (nearlyInstances.isEmpty()) {
            return rand(instances);
        }
        return rand(nearlyInstances);
    }

    private Mono<Response<ServiceInstance>> rand(List<ServiceInstance> instances) {
        ServiceInstance instance = instances.get(random.nextInt(instances.size()));
        return Mono.just(new DefaultResponse(instance));
    }
}